﻿using Ndns.Packet.Helpers;
using System;
using System.Collections.Generic;
using System.IO;

namespace Ndns.Packet
{
    public class DnsPacket : IEquatable<DnsPacket>
    {

        /// <summary>
        /// 表示DNS数据包的编号。
        /// </summary>
        public ushort TransactionID { get; set; }

        #region Flags相关

        /// <summary>
        /// 表示DNS数据包的特征标志。
        /// </summary>
        public DnsPacketFlag Flag
        {
            get { return _flag; }
            set { _flag = value; }
        }
        private DnsPacketFlag _flag = DnsPacketFlag.RecursionDesired;


        public bool IsResponsePacket
        {
            get { return (_flag & DnsPacketFlag.ResponsePacket) > 0; }
            set
            {
                _flag = (_flag & (DnsPacketFlag)(DnsPacketFlag.MaxValue - DnsPacketFlag.ResponsePacket));
                if (value) _flag = _flag | DnsPacketFlag.ResponsePacket;
            }
        }

        public bool Truncated
        {
            get { return (_flag & DnsPacketFlag.Truncated) > 0; }
            set
            {
                _flag = (_flag & (DnsPacketFlag)(DnsPacketFlag.MaxValue - DnsPacketFlag.Truncated));
                if (value) _flag = _flag | DnsPacketFlag.Truncated;
            }
        }

        public bool RecursionDesired
        {
            get { return (_flag & DnsPacketFlag.RecursionDesired) > 0; }
            set
            {
                _flag = (_flag & (DnsPacketFlag)(DnsPacketFlag.MaxValue - DnsPacketFlag.RecursionDesired));
                if (value) _flag = _flag | DnsPacketFlag.RecursionDesired;
            }
        }

        public bool RecursionAvailable
        {
            get { return (_flag & DnsPacketFlag.RecursionAvailable) > 0; }
            set
            {
                _flag = (_flag & (DnsPacketFlag)(DnsPacketFlag.MaxValue - DnsPacketFlag.RecursionAvailable));
                if (value) _flag = _flag | DnsPacketFlag.RecursionAvailable;
            }
        }

        public bool AnswerAuthenticated
        {
            get { return (_flag & DnsPacketFlag.AnswerAuthenticated) > 0; }
            set
            {
                _flag = (_flag & (DnsPacketFlag)(DnsPacketFlag.MaxValue - DnsPacketFlag.AnswerAuthenticated));
                if (value) _flag = _flag | DnsPacketFlag.AnswerAuthenticated;
            }
        }

        public bool AcceptNonAuthenticatedData
        {
            get { return (_flag & DnsPacketFlag.AcceptNonAuthenticatedData) > 0; }
            set
            {
                _flag = (_flag & (DnsPacketFlag)(DnsPacketFlag.MaxValue - DnsPacketFlag.AcceptNonAuthenticatedData));
                if (value) _flag = _flag | DnsPacketFlag.AcceptNonAuthenticatedData;
            }
        }

        /// <summary>
        /// 获取/设置Flag中的操作码字段。
        /// </summary>
        public DnsOpCode OpCode
        {
            get { return (DnsOpCode)(((ushort)_flag & 0x7800) >> 11); }
            set { _flag = (DnsPacketFlag)((ushort)_flag & 0x87ff | (((byte)value & 0xf) << 11)); }
        }

        /// <summary>
        /// 获取/设置Flag中的回复码字段。
        /// </summary>
        public DnsReplyCode ReplyCode
        {
            get { return (DnsReplyCode)((ushort)_flag & 0xf); }
            set { _flag = (DnsPacketFlag)((ushort)_flag & 0xfff0 | ((byte)value & 0xf)); }
        }

        #endregion

        #region 记录

        /// <summary>
        /// 数据包中包含的查询项。
        /// </summary>
        public IList<DnsQuery> Queries => _queries;

        /// <summary>
        /// 数据包中包含的应答记录。
        /// </summary>
        public IList<DnsResourceRecord> AnswerRRs => _answerRRs;

        /// <summary>
        /// 数据包中包含的权威命名服务器的记录。
        /// </summary>
        public IList<DnsResourceRecord> AuthoritativeNSs => _authoritativeNSs;

        /// <summary>
        /// 数据包中包含的附加记录。
        /// </summary>
        public IList<DnsResourceRecord> AdditionalRRs => _additionalRRs;

        private List<DnsQuery> _queries = new List<DnsQuery>();
        private List<DnsResourceRecord> _answerRRs = new List<DnsResourceRecord>();
        private List<DnsResourceRecord> _authoritativeNSs = new List<DnsResourceRecord>();
        private List<DnsResourceRecord> _additionalRRs = new List<DnsResourceRecord>();
        #endregion

        #region 实例构造

        /// <summary>
        /// 以默认属性构造DnsPacket实例。
        /// </summary>
        public DnsPacket()
        {

        }

        /// <summary>
        /// 构造DnsPacket实例，并设置TransactionID和DNS标志。
        /// </summary>
        /// <param name="transactionID">要设置的TransactionID属性。</param>
        /// <param name="flag">DNS标志</param>
        public DnsPacket(ushort transactionID, DnsPacketFlag flag)
        {
            TransactionID = transactionID;
            _flag = flag;
        }

        public static DnsPacket FromPacket(byte[] bytes)
        {
            using (var ms = new MemoryStream(bytes))
            {
                return FromPacket(ms);
            }
        }
        public static DnsPacket FromPacket(byte[] bytes, int index, int count)
        {
            using (var ms = new MemoryStream(bytes, index, count))
            {
                return FromPacket(ms);
            }
        }
        public static DnsPacket FromPacket(Stream s)
        {
            DnsPacket packet = new DnsPacket();
            long packetStartIndex = s.Position;
            //获取DNS头部
            int trnasactionID = s.ReadNetworkUInt16();
            int flag = s.ReadNetworkUInt16();
            if (trnasactionID == -1 || flag == -1)
                throw new IOException();
            packet.TransactionID = (ushort)trnasactionID;
            packet._flag = (DnsPacketFlag)flag;
            //获取各部分数量
            int cQueries, cAnswerRRs, cAuthoritativeNSs, cAdditionalRRs;
            cQueries = s.ReadNetworkUInt16();
            cAnswerRRs = s.ReadNetworkUInt16();
            cAuthoritativeNSs = s.ReadNetworkUInt16();
            cAdditionalRRs = s.ReadNetworkUInt16();
            if (cQueries == -1 || cAnswerRRs == -1 || cAuthoritativeNSs == -1 || cAdditionalRRs == -1)
                throw new IOException();
            //获取各个部分
            for (int i = 1; i <= cQueries; i++)
                try
                {
                    packet._queries.Add(DnsQuery.FromPacket(s, packetStartIndex));
                }
                catch (PacketEndOfStreamException ex)
                {
                    if (packet.Truncated)
                        return packet;
                    else
                        throw new IOException("数据包超出流范围。", ex);
                }

            for (int i = 1; i <= cAnswerRRs; i++)
                try
                {
                    packet._answerRRs.Add(DnsResourceRecord.FromPacket(s, packetStartIndex));
                }
                catch (PacketEndOfStreamException ex)
                {
                    if (packet.Truncated)
                        return packet;
                    else
                        throw new IOException("数据包超出流范围。", ex);
                }

            for (int i = 1; i <= cAuthoritativeNSs; i++)
                try
                {
                    packet._authoritativeNSs.Add(DnsResourceRecord.FromPacket(s, packetStartIndex));
                }
                catch (PacketEndOfStreamException ex)
                {
                    if (packet.Truncated)
                        return packet;
                    else
                        throw new IOException("数据包超出流范围。", ex);
                }
            for (int i = 1; i <= cAdditionalRRs; i++)
                try
                {
                    packet._additionalRRs.Add(DnsResourceRecord.FromPacket(s, packetStartIndex));
                }
                catch (PacketEndOfStreamException ex)
                {
                    if (packet.Truncated)
                        return packet;
                    else
                        throw new IOException("数据包超出流范围。", ex);
                }

            return packet;
        }

        #endregion

        #region 构造相关实例

        public DnsPacket CloneAsResponse(ushort transactionID, DnsReplyCode replyCode)
        {
            var packet = new DnsPacket(transactionID, this.Flag);
            packet.IsResponsePacket = true;
            packet.ReplyCode = replyCode;
            foreach (var item in _queries) packet.Queries.Add(item.Clone());
            return packet;
        }
        public DnsPacket CloneAsResponse(DnsReplyCode replyCode)
        {
            return CloneAsResponse(TransactionID, replyCode);
        }

        #endregion

        #region 获取信息

        public byte[] GetBytes()
        {
            return GetBytes(false, true);
        }

        public byte[] GetBytes(bool allowTruncation)
        {
            return GetBytes(allowTruncation, true);
        }

        public byte[] GetBytes(bool allowTruncation, bool enableCompression)
        {
            byte[] buff;
            //获得字节
            using (var ms = new MemoryStream(512))
            {
                if (enableCompression)
                    CopyCompressedTo(ms);
                else
                    CopyTo(ms);
                buff = ms.ToArray();
            }
            //处理截断
            if (allowTruncation)
            {
                if (buff.Length > 512)
                {
                    var truncatedBuff = new byte[512];
                    Buffer.BlockCopy(buff, 0, truncatedBuff, 0, 512);
                    buff = truncatedBuff;
                }
            }
            return buff;
        }

        /// <summary>
        /// 将此实例表示的数据包写入到流中。
        /// </summary>
        public int CopyTo(Stream s)
        {
            int count;

            //DNS数据包头
            s.WriteNetworkUInt16(TransactionID);
            s.WriteNetworkUInt16((ushort)_flag);

            //各个记录的数量
            s.WriteNetworkUInt16((ushort)_queries.Count);
            s.WriteNetworkUInt16((ushort)_answerRRs.Count);
            s.WriteNetworkUInt16((ushort)_authoritativeNSs.Count);
            s.WriteNetworkUInt16((ushort)_additionalRRs.Count);

            count = 12;

            //各条记录
            foreach (var e in _queries)
                count += e.CopyTo(s);
            foreach (var e in _answerRRs)
                count += e.CopyTo(s);
            foreach (var e in _authoritativeNSs)
                count += e.CopyTo(s);
            foreach (var e in _additionalRRs)
                count += e.CopyTo(s);

            return count;
        }

        public int CopyCompressedTo(Stream s)
        {
            int count;

            //字典
            var symbols = new Dictionary<string, DnsHostNamePointer>();
            int offset;
            var buffer = new byte[128];

            //DNS数据包头
            s.WriteNetworkUInt16(TransactionID);
            s.WriteNetworkUInt16((ushort)_flag);

            //各个记录的数量
            s.WriteNetworkUInt16((ushort)_queries.Count);
            s.WriteNetworkUInt16((ushort)_answerRRs.Count);
            s.WriteNetworkUInt16((ushort)_authoritativeNSs.Count);
            s.WriteNetworkUInt16((ushort)_additionalRRs.Count);

            //设置偏移量
            offset = 6 * sizeof(ushort);
            count = 12;

            //各条记录
            foreach (var e in _queries)
                count += e.CopyCompressedTo(s, symbols, ref offset);
            foreach (var e in _answerRRs)
                count += e.CopyCompressedTo(s, symbols, ref offset, buffer);
            foreach (var e in _authoritativeNSs)
                count += e.CopyCompressedTo(s, symbols, ref offset, buffer);
            foreach (var e in _additionalRRs)
                count += e.CopyCompressedTo(s, symbols, ref offset, buffer);

            return count;
        }
        #endregion

        #region Clone,Equals,HashCode实现

        public DnsPacket Clone()
        {
            DnsPacket packet = new DnsPacket(TransactionID, _flag);
            foreach (var e in _queries)
                packet._queries.Add(e.Clone());
            foreach (var e in _answerRRs)
                packet._answerRRs.Add(e.Clone());
            foreach (var e in _authoritativeNSs)
                packet._authoritativeNSs.Add(e.Clone());
            foreach (var e in _additionalRRs)
                packet._additionalRRs.Add(e.Clone());
            return packet;
        }

        public static bool Equals(DnsPacket a, DnsPacket b)
        {
            if (ReferenceEquals(a, b)) return true;
            if ((object)a == null || (object)b == null) return false;

            if (a.TransactionID != b.TransactionID) return false;
            if (a._flag != b._flag) return false;

            if (a._queries.Count != b._queries.Count) return false;
            if (a._answerRRs.Count != b._answerRRs.Count) return false;
            if (a._authoritativeNSs.Count != b._authoritativeNSs.Count) return false;
            if (a._additionalRRs.Count != b._additionalRRs.Count) return false;

            for (int i = 0; i < a._queries.Count; i++)
                if (!a._queries[i].Equals(b._queries[i]))
                    return false;
            for (int i = 0; i < a._answerRRs.Count; i++)
                if (!a._answerRRs[i].Equals(b._answerRRs[i]))
                    return false;
            for (int i = 0; i < a._authoritativeNSs.Count; i++)
                if (!a._authoritativeNSs[i].Equals(b._authoritativeNSs[i]))
                    return false;
            for (int i = 0; i < a._additionalRRs.Count; i++)
                if (!a._additionalRRs[i].Equals(b._additionalRRs[i]))
                    return false;

            return true;
        }

        public static bool operator ==(DnsPacket a, DnsPacket b)
        {
            return Equals(a, b);
        }

        public static bool operator !=(DnsPacket a, DnsPacket b)
        {
            return !Equals(a, b);
        }

        public bool Equals(DnsPacket other)
        {
            return Equals(this, other);
        }

        public override bool Equals(object obj)
        {
            DnsPacket other = obj as DnsPacket;
            if (other != null)
                return Equals(this, other);

            return base.Equals(obj);
        }

        public override int GetHashCode()
        {
            int total = 17;
            total = total * 37 + TransactionID.GetHashCode();
            total = total * 37 + _flag.GetHashCode();
            total = total * 37 + HashCodeHelper.GetHashCode(_answerRRs);
            total = total * 37 + HashCodeHelper.GetHashCode(_answerRRs);
            total = total * 37 + HashCodeHelper.GetHashCode(_authoritativeNSs);
            total = total * 37 + HashCodeHelper.GetHashCode(_additionalRRs);
            return total;
        }
        #endregion

    }
}
